use std::collections::VecDeque;

#[derive(Debug)]
enum Operation {
    Squared,
    Times(u64),
    Plus(u64),
}

#[derive(Debug)]
struct Monkey {
    items: VecDeque<u64>,
    operation: Operation,
    mod_number: u64,
    true_throw_target: usize,
    false_throw_target: usize,
    inspect_count: usize,
}

impl From<&str> for Monkey {
    /*
        Monkey 0:
            Starting items: 53, 89, 62, 57, 74, 51, 83, 97
            Operation: new = old * 3
            Test: divisible by 13
                If true: throw to monkey 1
                If false: throw to monkey 5
    */
    fn from(value: &str) -> Self {
        let lines = value.lines().map(|line| line.trim()).collect::<Vec<_>>();

        let colon_pos = lines[1].find(':').unwrap();
        let items = lines[1][colon_pos+2..]
            .split(", ")
            .map(|item| item.parse::<u64>().unwrap())
            .collect::<VecDeque<_>>();
        
        let equal_pos = lines[2].find('=').unwrap();
        let operation = match &lines[2][equal_pos+2..] {
            "old * old" => Operation::Squared,
            _ => {
                let constant = lines[2][equal_pos+8..].parse::<u64>().unwrap();
                match &lines[2][equal_pos+6..equal_pos+7] {
                    "+" => Operation::Plus(constant),
                    "*" => Operation::Times(constant),
                    _ => panic!("Unsupport oprtation: {}", &lines[2][equal_pos+2..]),
                }
            }
        };

        let mod_number = lines[3][19..].parse::<u64>().unwrap();

        let true_throw_target = lines[4][25..].parse::<usize>().unwrap();
        let false_throw_target = lines[5][26..].parse::<usize>().unwrap();

        Monkey {
            items, operation, mod_number, true_throw_target, false_throw_target, inspect_count: 0
        }
    }
}

fn parse_input() -> Vec<Monkey> {
    include_str!("../input.txt")
        .split("\r\n\r\n")
        .map(Into::<Monkey>::into)
        .collect()
}

fn part1() {
    let mut monkeys = parse_input();
    let monkey_size = monkeys.len();

    for _ in 0..20 {
        for monkey_id in 0..monkey_size {
            while let Some(item) = monkeys[monkey_id].items.pop_front() {
                monkeys[monkey_id].inspect_count += 1;
    
                let after_oper = match monkeys[monkey_id].operation {
                    Operation::Squared => item * item,
                    Operation::Times(num) => item * num,
                    Operation::Plus(num) => item + num,
                };

                let divide_3 = after_oper / 3;
                let throw_monkey_id = if divide_3 % monkeys[monkey_id].mod_number == 0 {
                    monkeys[monkey_id].true_throw_target
                } else {
                    monkeys[monkey_id].false_throw_target
                };

                monkeys[throw_monkey_id].items.push_back(divide_3);
            }
        }
    }

    monkeys.sort_by_key(|monkey| monkey.inspect_count);
    
    println!(
        "Part 1: {}", 
        monkeys[monkey_size - 1].inspect_count * monkeys[monkey_size - 2].inspect_count
    );

}

fn part2() {
    let mut monkeys = parse_input();
    let monkey_size = monkeys.len();

    let mod_number = monkeys.iter().map(|m: &Monkey| m.mod_number).product::<u64>();

    for _ in 0..10000 {
        for monkey_id in 0..monkey_size {
            while let Some(item) = monkeys[monkey_id].items.pop_front() {
                monkeys[monkey_id].inspect_count += 1;
    
                let after_oper = match monkeys[monkey_id].operation {
                    Operation::Squared => item * item,
                    Operation::Times(num) => item * num,
                    Operation::Plus(num) => item + num,
                };

                let after_oper = after_oper % mod_number;

                let throw_monkey_id = if after_oper % monkeys[monkey_id].mod_number == 0 {
                    monkeys[monkey_id].true_throw_target
                } else {
                    monkeys[monkey_id].false_throw_target
                };

                monkeys[throw_monkey_id].items.push_back(after_oper);
            }
        }
    }

    monkeys.sort_by_key(|monkey| monkey.inspect_count);
    
    println!(
        "Part 2: {}", 
        monkeys[monkey_size - 1].inspect_count * monkeys[monkey_size - 2].inspect_count
    );


}

fn main() {
    part1();
    part2();
}
